import os
import pandas as pd
from zipfile import ZipFile
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.metrics.pairwise import cosine_similarity


names_path = 'names.xlsx'
homework_path = 'homework1'
dist_homework = 'tmp'

if not os.path.exists(dist_homework):
    os.makedirs(dist_homework)

df = pd.read_excel(names_path)

names = df['姓名'].to_list()
stunos = df['学号'].to_list()


def unzip_homework(file_name, stuno, name, dist='tmp'):
    with ZipFile(file_name, 'r') as zObject:
        for name1 in zObject.namelist():
            if name1.endswith('ex1.py') or name1.endswith('ex2.py') or name1.endswith('ex3.py'):
                zObject.extract(name1, os.path.join(dist, f'{stuno}{name}'))
        zObject.close()

for stuno, name in zip(stunos, names):
    file_name = f'{stuno}{name}.zip'
    file_name1 = f'{stuno} {name}.zip'
    if os.path.exists(os.path.join(homework_path, file_name)):
        unzip_homework(os.path.join(homework_path, file_name), stuno, name)
    if os.path.exists(os.path.join(homework_path, file_name1)):
        unzip_homework(os.path.join(homework_path, file_name1), stuno, name)

contents = {}
for root, dirs, files in os.walk(dist_homework):
    # print(root)
    if files:
        result = ''
        for file in sorted(files):
            # print(file)
            with open(os.path.join(root, file), 'r', encoding='utf8') as f:
                result += f.read()
        contents[root.split('\\')[1]] = result
        # print(result)

# print(contents.keys())


def calculate_tfidf_cosine_similarity(text1, text2):
    vectorizer = TfidfVectorizer()
    corpus = [text1, text2]
    vectors = vectorizer.fit_transform(corpus)
    similarity = cosine_similarity(vectors)
    return similarity[0][1]


for name1, text1 in contents.items():
    for name2, text2 in contents.items():
        sim = calculate_tfidf_cosine_similarity(text1, text2)
        print(f'{name1}和{name2}的相似度：{sim}')

